import sys
import os
import subprocess
import torch
import whisper
import tempfile
import shutil
from datetime import timedelta
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                               QHBoxLayout, QPushButton, QLabel, QFileDialog,
                               QComboBox, QSpinBox, QLineEdit, QProgressBar,
                               QTextEdit, QMessageBox, QTabWidget, QGroupBox,
                               QFormLayout, QCheckBox)
from PySide6.QtCore import Qt, QThread, Signal, QSettings
from PySide6.QtGui import QIcon, QPixmap


class SubtitleGeneratorThread(QThread):
    """后台处理视频和生成字幕的线程"""
    progress_updated = Signal(str, int)
    task_completed = Signal(bool, str)
    log_message = Signal(str)

    def __init__(self, video_path, output_path, model_size, font_size,
                 font_color, position, keep_temp):
        super().__init__()
        self.video_path = video_path
        self.output_path = output_path
        self.model_size = model_size
        self.font_size = font_size
        self.font_color = font_color
        self.position = position
        self.keep_temp = keep_temp
        self.temp_dir = None

    def run(self):
        try:
            self.log_message.emit("开始处理视频...")
            self.progress_updated.emit("准备中...", 5)

            # 创建临时目录
            self.temp_dir = tempfile.mkdtemp(prefix="video_subtitle_")
            self.log_message.emit(f"创建临时目录: {self.temp_dir}")

            # 1. 提取音频
            self.progress_updated.emit("正在提取音频...", 10)
            audio_path = os.path.join(self.temp_dir, "temp_audio.wav")
            if not self._extract_audio(audio_path):
                raise Exception("音频提取失败")
            self.progress_updated.emit("音频提取完成", 25)

            # 2. 生成字幕
            self.progress_updated.emit(f"正在加载{self.model_size}模型...", 30)
            self.log_message.emit(f"使用设备: {'GPU' if torch.cuda.is_available() else 'CPU'}")

            subtitle_path = os.path.join(self.temp_dir, "subtitles.vtt")
            if not self._generate_subtitles(audio_path, subtitle_path):
                raise Exception("字幕生成失败")
            self.progress_updated.emit("字幕生成完成", 70)

            # 3. 合并字幕到视频
            self.progress_updated.emit("正在合并字幕到视频...", 75)
            if not self._merge_subtitles(subtitle_path):
                raise Exception("字幕合并失败")
            self.progress_updated.emit("字幕合并完成", 95)

            # 清理临时文件
            if not self.keep_temp and self.temp_dir:
                self.log_message.emit("清理临时文件...")
                shutil.rmtree(self.temp_dir, ignore_errors=True)
                self.temp_dir = None

            self.progress_updated.emit("处理完成", 100)
            self.task_completed.emit(True, f"视频处理完成!\n输出文件: {self.output_path}")
        except Exception as e:
            self.log_message.emit(f"错误: {str(e)}")
            self.progress_updated.emit("处理失败", 100)
            self.task_completed.emit(False, f"处理失败: {str(e)}")

            # 清理临时文件
            if self.temp_dir:
                shutil.rmtree(self.temp_dir, ignore_errors=True)
                self.temp_dir = None

    def _extract_audio(self, audio_path):
        """从视频中提取音频"""
        self.log_message.emit(f"从视频中提取音频: {self.video_path}")
        try:
            subprocess.run([
                'ffmpeg',
                '-i', self.video_path,
                '-vn',  # 无视频
                '-acodec', 'pcm_s16le',  # 音频编码
                audio_path
            ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            self.log_message.emit(f"音频提取完成，保存至: {audio_path}")
            return True
        except subprocess.CalledProcessError as e:
            self.log_message.emit(f"音频提取失败: {e.stderr.decode()}")
            return False

    def _generate_subtitles(self, audio_path, subtitle_path):
        """使用Whisper生成字幕"""
        self.log_message.emit(f"使用Whisper模型({self.model_size})生成字幕...")
        try:
            # 加载模型
            model = whisper.load_model(self.model_size)

            # 转录音频
            result = model.transcribe(audio_path)

            # 保存为VTT格式
            with open(subtitle_path, 'w', encoding='utf-8') as f:
                f.write("WEBVTT\n\n")
                for segment in result['segments']:
                    start_time = timedelta(seconds=segment['start'])
                    end_time = timedelta(seconds=segment['end'])
                    text = segment['text'].strip()

                    # 格式化时间
                    start_str = self._format_time(start_time)
                    end_str = self._format_time(end_time)

                    f.write(f"{start_str} --> {end_str}\n{text}\n\n")

            self.log_message.emit(f"字幕生成完成，保存至: {subtitle_path}")
            return True
        except Exception as e:
            self.log_message.emit(f"字幕生成失败: {str(e)}")
            return False

    def _merge_subtitles(self, subtitle_path):
        """将字幕合并到视频中"""
        self.log_message.emit(f"将字幕合并到视频中...")

        # 确定字幕位置参数
        if self.position == 'bottom':
            position_args = 'y=(h-text_h-20)'
        elif self.position == 'top':
            position_args = 'y=20'
        else:
            # 尝试解析自定义位置
            try:
                if '%' in self.position:
                    percent = float(self.position.strip('%')) / 100
                    position_args = f'y=h*{percent}-text_h/2'
                else:
                    position_args = f'y={self.position}'
            except:
                self.log_message.emit("无效的位置参数，使用默认底部位置")
                position_args = 'y=(h-text_h-20)'

        try:
            # 构建ffmpeg命令
            cmd = [
                'ffmpeg',
                '-i', self.video_path,
                '-vf',
                f"subtitles={subtitle_path}:force_style='FontSize={self.font_size},PrimaryColour=&H{self._color_to_bgr(self.font_color)}'",
                '-c:a', 'copy',  # 保持音频不变
                '-y',  # 覆盖已存在的文件
                self.output_path
            ]

            # 执行命令
            subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            self.log_message.emit(f"字幕合并完成，保存至: {self.output_path}")
            return True
        except subprocess.CalledProcessError as e:
            self.log_message.emit(f"字幕合并失败: {e.stderr.decode()}")
            return False

    def _format_time(self, time_obj):
        """格式化时间为00:00:00.000格式"""
        time_str = str(time_obj).replace(',', '.')

        # 如果时间格式不包含毫秒，添加毫秒部分
        if '.' not in time_str:
            time_str += '.000'

        # 确保时间格式为HH:MM:SS.XXX
        parts = time_str.split(':')
        if len(parts) == 2:
            time_str = '00:' + time_str

        return time_str

    def _color_to_bgr(self, color):
        """将颜色名称或十六进制值转换为BGR格式（用于FFmpeg）"""
        # 常见颜色映射
        color_map = {
            'white': 'FFFFFF',
            'black': '000000',
            'red': 'FF0000',
            'green': '00FF00',
            'blue': '0000FF',
            'yellow': 'FFFF00',
            'cyan': '00FFFF',
            'magenta': 'FF00FF'
        }

        # 如果是颜色名称，转换为十六进制
        if color.lower() in color_map:
            hex_color = color_map[color.lower()]
        else:
            # 去除可能的#前缀
            hex_color = color.lstrip('#')

        # 确保是6位十六进制
        if len(hex_color) != 6:
            self.log_message.emit(f"无效的颜色值: {color}，使用白色作为默认值")
            hex_color = 'FFFFFF'

        return hex_color


class VideoSubtitleGenerator(QMainWindow):
    """视频字幕生成工具主窗口"""

    def __init__(self):
        super().__init__()

        # 设置应用程序设置
        self.settings = QSettings("VideoSubtitleGenerator", "PySide6")

        # 设置窗口属性
        self.setWindowTitle("视频字幕生成工具")
        self.setMinimumSize(800, 600)

        # 创建中央部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # 创建选项卡控件
        self.tab_widget = QTabWidget()

        # 添加主界面选项卡
        self.main_tab = QWidget()
        self.tab_widget.addTab(self.main_tab, "主界面")
        self._setup_main_tab()

        # 添加日志选项卡
        self.log_tab = QWidget()
        self.tab_widget.addTab(self.log_tab, "处理日志")
        self._setup_log_tab()

        main_layout.addWidget(self.tab_widget)

        # 初始化变量
        self.worker_thread = None

        # 恢复上次的设置
        self._restore_settings()

    def _setup_main_tab(self):
        """设置主界面选项卡"""
        layout = QVBoxLayout(self.main_tab)

        # 输入文件部分
        input_group = QGroupBox("输入文件")
        input_layout = QVBoxLayout(input_group)

        file_layout = QHBoxLayout()
        self.input_file_label = QLabel("未选择文件")
        self.browse_button = QPushButton("浏览...")
        self.browse_button.clicked.connect(self._browse_file)

        file_layout.addWidget(self.input_file_label, 1)
        file_layout.addWidget(self.browse_button)

        input_layout.addLayout(file_layout)

        # 输出文件部分
        output_group = QGroupBox("输出文件")
        output_layout = QVBoxLayout(output_group)

        output_file_layout = QHBoxLayout()
        self.output_file_label = QLabel("未设置")
        self.output_browse_button = QPushButton("浏览...")
        self.output_browse_button.clicked.connect(self._browse_output_file)

        output_file_layout.addWidget(self.output_file_label, 1)
        output_file_layout.addWidget(self.output_browse_button)

        output_layout.addLayout(output_file_layout)

        # 字幕设置部分
        settings_group = QGroupBox("字幕设置")
        settings_layout = QFormLayout(settings_group)

        # 模型选择
        self.model_combo = QComboBox()
        self.model_combo.addItems(["tiny", "base", "small", "medium", "large"])
        self.model_combo.setToolTip("模型越大，识别越准确，但处理速度越慢")

        # 字体大小
        self.font_size_spin = QSpinBox()
        self.font_size_spin.setRange(12, 72)
        self.font_size_spin.setValue(24)

        # 字体颜色
        self.font_color_edit = QLineEdit("white")
        self.font_color_edit.setToolTip("可以使用颜色名称(如white, red)或十六进制值(如#FF0000)")

        # 字幕位置
        self.position_combo = QComboBox()
        self.position_combo.addItems(["bottom", "top", "自定义"])
        self.position_combo.currentTextChanged.connect(self._update_position_ui)

        self.position_edit = QLineEdit("50%")
        self.position_edit.setEnabled(False)

        position_layout = QHBoxLayout()
        position_layout.addWidget(self.position_combo)
        position_layout.addWidget(self.position_edit, 1)

        # 保留临时文件
        self.keep_temp_checkbox = QCheckBox("保留临时文件")

        # 添加到布局
        settings_layout.addRow("Whisper模型:", self.model_combo)
        settings_layout.addRow("字幕字体大小:", self.font_size_spin)
        settings_layout.addRow("字幕字体颜色:", self.font_color_edit)
        settings_layout.addRow("字幕位置:", position_layout)
        settings_layout.addRow(self.keep_temp_checkbox)

        # 进度条
        self.progress_label = QLabel("准备就绪")
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)

        # 开始按钮
        self.start_button = QPushButton("开始处理")
        self.start_button.setMinimumHeight(40)
        self.start_button.setStyleSheet("font-size: 14px; font-weight: bold;")
        self.start_button.clicked.connect(self._start_processing)

        # 添加所有部分到主布局
        layout.addWidget(input_group)
        layout.addWidget(output_group)
        layout.addWidget(settings_group)
        layout.addWidget(self.progress_label)
        layout.addWidget(self.progress_bar)
        layout.addWidget(self.start_button)

        layout.addStretch(1)

    def _setup_log_tab(self):
        """设置日志选项卡"""
        layout = QVBoxLayout(self.log_tab)

        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setStyleSheet("font-family: Consolas, monospace;")

        buttons_layout = QHBoxLayout()
        self.clear_log_button = QPushButton("清空日志")
        self.clear_log_button.clicked.connect(self._clear_log)

        self.copy_log_button = QPushButton("复制日志")
        self.copy_log_button.clicked.connect(self._copy_log)

        buttons_layout.addWidget(self.clear_log_button)
        buttons_layout.addWidget(self.copy_log_button)

        layout.addWidget(self.log_text)
        layout.addLayout(buttons_layout)

    def _browse_file(self):
        """浏览并选择视频文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", "", "视频文件 (*.mp4 *.mkv *.avi *.mov *.wmv);;所有文件 (*)"
        )

        if file_path:
            self.input_file_label.setText(file_path)

            # 自动设置输出文件路径
            base, ext = os.path.splitext(file_path)
            self.output_file_label.setText(f"{base}_subtitles.mp4")

            # 保存设置
            self.settings.setValue("last_input_dir", os.path.dirname(file_path))

    def _browse_output_file(self):
        """浏览并选择输出文件路径"""
        if self.input_file_label.text() and os.path.exists(self.input_file_label.text()):
            default_dir = os.path.dirname(self.input_file_label.text())
        else:
            default_dir = self.settings.value("last_output_dir", "")

        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存输出文件", default_dir, "视频文件 (*.mp4);;所有文件 (*)"
        )

        if file_path:
            self.output_file_label.setText(file_path)
            self.settings.setValue("last_output_dir", os.path.dirname(file_path))

    def _update_position_ui(self, position):
        """根据位置选择更新UI状态"""
        self.position_edit.setEnabled(position == "自定义")

    def _start_processing(self):
        """开始处理视频"""
        # 检查输入文件
        input_file = self.input_file_label.text()
        if not input_file or input_file == "未选择文件":
            QMessageBox.warning(self, "警告", "请选择输入视频文件")
            return

        if not os.path.exists(input_file):
            QMessageBox.warning(self, "警告", f"输入文件不存在:\n{input_file}")
            return

        # 检查输出文件
        output_file = self.output_file_label.text()
        if not output_file or output_file == "未设置":
            base, ext = os.path.splitext(input_file)
            output_file = f"{base}_subtitles.mp4"
            self.output_file_label.setText(output_file)

        # 检查输出目录是否存在
        output_dir = os.path.dirname(output_file)
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
            except:
                QMessageBox.warning(self, "警告", f"无法创建输出目录:\n{output_dir}")
                return

        # 确认是否覆盖已存在的文件
        if os.path.exists(output_file):
            reply = QMessageBox.question(
                self, "确认", f"输出文件已存在:\n{output_file}\n是否覆盖?",
                QMessageBox.Yes | QMessageBox.No, QMessageBox.No
            )
            if reply != QMessageBox.Yes:
                return

        # 获取设置
        model_size = self.model_combo.currentText()
        font_size = self.font_size_spin.value()
        font_color = self.font_color_edit.text()

        position = self.position_combo.currentText()
        if position == "自定义":
            position = self.position_edit.text()

        keep_temp = self.keep_temp_checkbox.isChecked()

        # 保存设置
        self._save_settings()

        # 切换到日志选项卡
        self.tab_widget.setCurrentIndex(1)

        # 清空日志
        self.log_text.clear()

        # 禁用UI元素
        self.start_button.setEnabled(False)
        self.browse_button.setEnabled(False)
        self.output_browse_button.setEnabled(False)

        # 创建并启动工作线程
        self.worker_thread = SubtitleGeneratorThread(
            input_file, output_file, model_size, font_size,
            font_color, position, keep_temp
        )
        self.worker_thread.progress_updated.connect(self._update_progress)
        self.worker_thread.task_completed.connect(self._processing_completed)
        self.worker_thread.log_message.connect(self._log_message)
        self.worker_thread.start()

    def _update_progress(self, message, value):
        """更新进度条和状态标签"""
        self.progress_label.setText(message)
        self.progress_bar.setValue(value)
        self._log_message(message)

    def _processing_completed(self, success, message):
        """处理完成后的回调函数"""
        # 重新启用UI元素
        self.start_button.setEnabled(True)
        self.browse_button.setEnabled(True)
        self.output_browse_button.setEnabled(True)

        # 显示消息框
        if success:
            QMessageBox.information(self, "成功", message)
        else:
            QMessageBox.critical(self, "失败", message)

    def _log_message(self, message):
        """添加日志消息"""
        self.log_text.append(message)
        # 滚动到底部
        self.log_text.moveCursor(self.log_text.textCursor().End)

    def _clear_log(self):
        """清空日志"""
        self.log_text.clear()

    def _copy_log(self):
        """复制日志到剪贴板"""
        self.log_text.selectAll()
        self.log_text.copy()
        self.log_text.moveCursor(self.log_text.textCursor().End)
        QMessageBox.information(self, "提示", "日志已复制到剪贴板")

    def _save_settings(self):
        """保存应用程序设置"""
        self.settings.setValue("model_size", self.model_combo.currentText())
        self.settings.setValue("font_size", self.font_size_spin.value())
        self.settings.setValue("font_color", self.font_color_edit.text())
        self.settings.setValue("position", self.position_combo.currentText())
        self.settings.setValue("custom_position", self.position_edit.text())
        self.settings.setValue("keep_temp", self.keep_temp_checkbox.isChecked())

    def _restore_settings(self):
        """恢复应用程序设置"""
        model_size = self.settings.value("model_size", "small")
        font_size = int(self.settings.value("font_size", 24))
        font_color = self.settings.value("font_color", "white")
        position = self.settings.value("position", "bottom")
        custom_position = self.settings.value("custom_position", "50%")
        keep_temp = self.settings.value("keep_temp", False) == "true"

        self.model_combo.setCurrentText(model_size)
        self.font_size_spin.setValue(font_size)
        self.font_color_edit.setText(font_color)
        self.position_combo.setCurrentText(position)
        self.position_edit.setText(custom_position)
        self.keep_temp_checkbox.setChecked(keep_temp)

        self._update_position_ui(position)


def main():
    # 确保中文显示正常
    import matplotlib
    matplotlib.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

    app = QApplication(sys.argv)

    #
    app.setStyle("Fusion")

    # 创建并显示窗口
    window = VideoSubtitleGenerator()
    window.show()

    sys.exit(app.exec())


if __name__ == "__main__":
    main()    